/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.runners.core.metrics;



import com.bff.gaia.unified.model.fnexecution.v1.UnifiedFnApi;

import com.bff.gaia.unified.sdk.annotations.Experimental;

import com.bff.gaia.unified.sdk.annotations.Experimental.Kind;

import com.bff.gaia.unified.sdk.metrics.MetricKey;

import com.bff.gaia.unified.sdk.metrics.MetricName;

import com.bff.gaia.unified.vendor.guava.com.google.common.cache.CacheBuilder;

import com.bff.gaia.unified.vendor.guava.com.google.common.cache.CacheLoader;

import com.bff.gaia.unified.vendor.guava.com.google.common.cache.LoadingCache;



import java.util.ArrayList;

import java.util.Collection;

import java.util.List;

import java.util.Map;



/** Translation utilities for metrics classes to/from Fn API. */

@Experimental(Kind.METRICS)

public abstract class MetricsTranslation {



  private MetricsTranslation() {}



  public static MetricUpdates metricUpdatesFromProto(

      String ptransformName, Collection<UnifiedFnApi.Metrics.User> userMetricUpdates) {

    List<MetricUpdates.MetricUpdate<Long>> counterUpdates = new ArrayList<>();

    List<MetricUpdates.MetricUpdate<DistributionData>> distributionUpdates = new ArrayList<>();

    List<MetricUpdates.MetricUpdate<GaugeData>> gaugeUpdates = new ArrayList<>();



    for (UnifiedFnApi.Metrics.User userMetricUpdate : userMetricUpdates) {

      MetricKey metricKey =

          MetricKey.create(ptransformName, metricNameFromProto(userMetricUpdate.getMetricName()));

      switch (userMetricUpdate.getDataCase()) {

        case COUNTER_DATA:

          counterUpdates.add(

              MetricUpdates.MetricUpdate.create(

                  metricKey, userMetricUpdate.getCounterData().getValue()));

          break;

        case DISTRIBUTION_DATA:

          distributionUpdates.add(

              MetricUpdates.MetricUpdate.create(

                  metricKey,

                  DistributionData.create(

                      userMetricUpdate.getDistributionData().getSum(),

                      userMetricUpdate.getDistributionData().getCount(),

                      userMetricUpdate.getDistributionData().getMin(),

                      userMetricUpdate.getDistributionData().getMax())));

          break;

        case GAUGE_DATA:

          gaugeUpdates.add(

              MetricUpdates.MetricUpdate.create(

                  metricKey, GaugeData.create(userMetricUpdate.getGaugeData().getValue())));

          break;

        case DATA_NOT_SET:

          continue;

      }

    }

    return MetricUpdates.create(counterUpdates, distributionUpdates, gaugeUpdates);

  }



  public static Map<String, Collection<UnifiedFnApi.Metrics.User>> metricUpdatesToProto(

      MetricUpdates metricUpdates) {

    LoadingCache<String, Collection<UnifiedFnApi.Metrics.User>> fnMetrics =

        CacheBuilder.newBuilder()

            .build(

                new CacheLoader<String, Collection<UnifiedFnApi.Metrics.User>>() {

                  @Override

                  public Collection<UnifiedFnApi.Metrics.User> load(String ptransformName) {

                    return new ArrayList<>();

                  }

                });



    for (MetricUpdates.MetricUpdate<Long> counterUpdate : metricUpdates.counterUpdates()) {

      fnMetrics

          .getUnchecked(counterUpdate.getKey().stepName())

          .add(

              UnifiedFnApi.Metrics.User.newBuilder()

                  .setMetricName(metricNameToProto(counterUpdate.getKey().metricName()))

                  .setCounterData(

                      UnifiedFnApi.Metrics.User.CounterData.newBuilder()

                          .setValue(counterUpdate.getUpdate()))

                  .build());

    }



    for (MetricUpdates.MetricUpdate<GaugeData> gaugeUpdate : metricUpdates.gaugeUpdates()) {

      fnMetrics

          .getUnchecked(gaugeUpdate.getKey().stepName())

          .add(

              UnifiedFnApi.Metrics.User.newBuilder()

                  .setMetricName(metricNameToProto(gaugeUpdate.getKey().metricName()))

                  .setGaugeData(

                      UnifiedFnApi.Metrics.User.GaugeData.newBuilder()

                          .setValue(gaugeUpdate.getUpdate().value()))

                  .build());

    }



    for (MetricUpdates.MetricUpdate<DistributionData> distributionUpdate :

        metricUpdates.distributionUpdates()) {

      fnMetrics

          .getUnchecked(distributionUpdate.getKey().stepName())

          .add(

              UnifiedFnApi.Metrics.User.newBuilder()

                  .setMetricName(metricNameToProto(distributionUpdate.getKey().metricName()))

                  .setDistributionData(

                      UnifiedFnApi.Metrics.User.DistributionData.newBuilder()

                          .setCount(distributionUpdate.getUpdate().count())

                          .setMax(distributionUpdate.getUpdate().max())

                          .setMin(distributionUpdate.getUpdate().min())

                          .setSum(distributionUpdate.getUpdate().sum()))

                  .build());

    }



    return fnMetrics.asMap();

  }



  public static MetricName metricNameFromProto(UnifiedFnApi.Metrics.User.MetricName protoMetricName) {

    return MetricName.named(protoMetricName.getNamespace(), protoMetricName.getName());

  }



  public static UnifiedFnApi.Metrics.User.MetricName metricNameToProto(MetricName metricName) {

    return UnifiedFnApi.Metrics.User.MetricName.newBuilder()

        .setNamespace(metricName.getNamespace())

        .setName(metricName.getName())

        .build();

  }

}